{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center>\n",
    "    <p style=\"text-align:center\">\n",
    "        <img alt=\"phoenix logo\" src=\"https://storage.googleapis.com/arize-phoenix-assets/assets/phoenix-logo-light.svg\" width=\"200\"/>\n",
    "        <br>\n",
    "        <a href=\"https://arize.com/docs/phoenix/\">Docs</a>\n",
    "        |\n",
    "        <a href=\"https://github.com/Arize-ai/phoenix\">GitHub</a>\n",
    "        |\n",
    "        <a href=\"https://arize-ai.slack.com/join/shared_invite/zt-2w57bhem8-hq24MB6u7yE_ZF_ilOYSBw#/shared-invite/email\">Community</a>\n",
    "    </p>\n",
    "</center>\n",
    "\n",
    "# Google GenAI SDK - Building a Sequential Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q google-genai arize-phoenix-otel openinference-instrumentation-google-genai"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connect to Arize Phoenix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "from google import genai\n",
    "\n",
    "from phoenix.otel import register\n",
    "\n",
    "if \"PHOENIX_API_KEY\" not in os.environ:\n",
    "    os.environ[\"PHOENIX_API_KEY\"] = getpass(\"🔑 Enter your Phoenix API key: \")\n",
    "\n",
    "if \"PHOENIX_COLLECTOR_ENDPOINT\" not in os.environ:\n",
    "    os.environ[\"PHOENIX_COLLECTOR_ENDPOINT\"] = getpass(\"🔑 Enter your Phoenix Collector Endpoint\")\n",
    "\n",
    "tracer_provider = register(auto_instrument=True, project_name=\"google-genai-sequential-agent\")\n",
    "tracer = tracer_provider.get_tracer(__name__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Authenticate with Google Vertex AI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!gcloud auth login"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a client using the Vertex AI API, you could also use the Google GenAI API instead here\n",
    "client = genai.Client(vertexai=True, project=\"<ADD YOUR GCP PROJECT ID>\", location=\"us-central1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sequential Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- 1. Define helper methods ---\n",
    "\n",
    "\n",
    "def get_model(model_name):\n",
    "    \"\"\"Instantiate and return a model.\"\"\"\n",
    "    return model_name\n",
    "\n",
    "\n",
    "@tracer.chain\n",
    "def research_topic(client, model, user_input):\n",
    "    \"\"\"Research a topic based on user input.\"\"\"\n",
    "    research_prompt = (\n",
    "        \"You are a research assistant.\\n\"\n",
    "        \"Provide a comprehensive overview of the following topic.\\n\"\n",
    "        \"Include key facts, historical context, and current relevance.\\n\"\n",
    "        \"Keep your response to 3-4 paragraphs.\\n\"\n",
    "        \"User Input: \"\n",
    "    )\n",
    "    research_response = client.models.generate_content(\n",
    "        model=model,\n",
    "        contents=research_prompt + user_input,\n",
    "    )\n",
    "    return research_response.text.strip()\n",
    "\n",
    "\n",
    "@tracer.chain\n",
    "def identify_key_points(client, model, research):\n",
    "    \"\"\"Extract key points from the research.\"\"\"\n",
    "    key_points_prompt = (\n",
    "        \"You are a research summarizer.\\n\"\n",
    "        \"Extract 5-7 key points from the following research.\\n\"\n",
    "        \"Format each point as a bullet point with a brief explanation.\\n\\n\"\n",
    "        f\"Research:\\n{research}\"\n",
    "    )\n",
    "    key_points_response = client.models.generate_content(model=model, contents=key_points_prompt)\n",
    "    return key_points_response.text.strip()\n",
    "\n",
    "\n",
    "@tracer.chain\n",
    "def generate_future_directions(client, model, research, key_points):\n",
    "    \"\"\"Generate future research directions based on the research and key points.\"\"\"\n",
    "    future_prompt = (\n",
    "        \"You are a research strategist.\\n\"\n",
    "        \"Based on the research and key points provided, suggest 3-4 promising future research directions.\\n\"\n",
    "        \"For each direction, explain its potential significance and impact.\\n\\n\"\n",
    "        f\"Research:\\n{research}\\n\\n\"\n",
    "        f\"Key Points:\\n{key_points}\\n\\n\"\n",
    "    )\n",
    "    future_response = client.models.generate_content(model=model, contents=future_prompt)\n",
    "    return future_response.text.strip()\n",
    "\n",
    "\n",
    "# --- 2. Main execution ---\n",
    "@tracer.agent()\n",
    "def run_agent(user_input):\n",
    "    # Instantiate the models\n",
    "    research_model = get_model(\"gemini-2.0-flash-001\")\n",
    "\n",
    "    # Step 1: Generate the initial research\n",
    "    initial_research = research_topic(client, research_model, user_input)\n",
    "\n",
    "    # Step 2: Identify key points from the research\n",
    "    extracted_key_points = identify_key_points(client, research_model, initial_research)\n",
    "\n",
    "    # Step 3: Generate future research directions\n",
    "    future_directions = generate_future_directions(\n",
    "        client, research_model, initial_research, extracted_key_points\n",
    "    )\n",
    "\n",
    "    # Display the results\n",
    "    print(\"=== Initial Research ===\\n\")\n",
    "    print(initial_research)\n",
    "    print(\"\\n=== Key Points ===\\n\")\n",
    "    print(extracted_key_points)\n",
    "    print(\"\\n=== Future Research Directions ===\\n\")\n",
    "    print(future_directions)\n",
    "\n",
    "    return {\n",
    "        \"research\": initial_research,\n",
    "        \"key_points\": extracted_key_points,\n",
    "        \"future_directions\": future_directions,\n",
    "    }\n",
    "\n",
    "\n",
    "run_agent(user_input=input(\"Please enter a topic you'd like to research: \"))"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
